import os
import urllib.request as request

import tensorflow as tf
import numpy as np

# 本例使用虹膜数据集来进行TensorFlow的高级机器学习API

IRIS_TRAINING = "resource/iris/iris_training.csv"
IRIS_TRAINING_URL = "http://download.tensorflow.org/data/iris_training.csv"

IRIS_TEST = "resource/iris/iris_test.csv"
IRIS_TEST_URL = "http://download.tensorflow.org/data/iris_test.csv"


def main():

    # 检查下载CSV数据，如果未下载则下载
    if not os.path.exists(os.path.dirname(IRIS_TRAINING)):
        os.makedirs(os.path.dirname(IRIS_TRAINING))
    if not os.path.exists(IRIS_TRAINING):
        raw = request.urlopen(IRIS_TRAINING_URL).read()
        with open(IRIS_TRAINING, 'wb+') as f:
            f.write(raw)

    if not os.path.exists(os.path.dirname(IRIS_TEST)):
        os.makedirs(os.path.dirname(IRIS_TEST))
    if not os.path.exists(IRIS_TEST):
        raw = request.urlopen(IRIS_TEST_URL).read()
        with open(IRIS_TEST, 'wb+') as f:
            f.write(raw)

    # 加载数据
    training_set = tf.contrib.learn.datasets.base.load_csv_with_header(
        filename=IRIS_TRAINING,  # CSV文件路径
        target_dtype=np.int,  # 目标值类型
        features_dtype=np.float32  # 特征值类型
    )
    test_set = tf.contrib.learn.datasets.base.load_csv_with_header(
        filename=IRIS_TEST,
        target_dtype=np.int,
        features_dtype=np.float32
    )

    # 测试输出
    print(test_set.data)  # 特征值，默认为出去最后一列的所有列
    print(test_set.target)  # 目标值，默认为最后一列

    # 构建深层神经网络分类器
    feature_columns = [tf.contrib.layers.real_valued_column("", dimension=4)]  # 设置4个连续值特征数据组成特征列集合
    classifier = tf.contrib.learn.DNNClassifier(
        feature_columns=feature_columns,  # 特征列集合
        hidden_units=[10, 20, 10],  # 三个 隐藏层，分别包含10,20和10个神经元。
        n_classes=3,  # 三个目标课程，代表三种虹膜物种。
        model_dir="../data/model/iris"  # 保存检查点数据的目录
    )  # 创建一个深层神经网络模型

    # 描述训练模型的输入
    def get_train_inputs():
        """
        取得训练模型的输入数据
        :return:
        """
        x = tf.constant(training_set.data)
        y = tf.constant(training_set.target)
        return x, y

    # 训练该模型，模型的状态保留在classifier
    classifier.fit(
        input_fn=get_train_inputs,  # 指定输入方法
        steps=2000  # 训练步数
    )
    # classifier.fit(x=training_set.data, y=training_set.target, steps=2000)  # 等同于上个语句

    # 定义测试模型的输入
    def get_test_inputs():
        x = tf.constant(test_set.data)  # 特征值数组
        y = tf.constant(test_set.target)  # 目标值数组
        return x, y

    # 执行模型并取得值，模型的evaluate方法
    t_dict = classifier.evaluate(
        input_fn=get_test_inputs,
        steps=1
    )

    # 输出相关值
    print(t_dict)
    accuracy_score = t_dict['accuracy']
    print("\nTest Accuracy: {0:f}\n".format(accuracy_score))

    # 定义要进行预测的输入数据
    def new_samples():
        return np.array(
            [
                [6.4, 3.2, 4.5, 1.5],  # 第1个
                [5.8, 3.1, 5.0, 1.7]   # 第2个，可根据情况自己添加
            ],
            dtype=np.float32
        )

    # 使用模型的predict方法来进行预测
    predictions = list(classifier.predict(input_fn=new_samples))

    # 输出预测结果
    print("New Samples, Class Predictions:    {}\n".format(predictions))

if __name__ == "__main__":
    main()
